import numpy
import scipy.optimize
from kuai.sim.simtool import *
from kuai.sim.engine import *


class EnergyFunction:
    def __init__(self, mol, ff, energy):
        self.mol = mol
        self.ff = ff
        self.energy = energy
        
    def __call__(self, x):
        self.mol.coords = x
        return self.energy.energy(self.mol, self.ff)

class GradientFunction:
    def __init__(self, mol, ff, energy):
        self.mol = mol
        self.ff = ff
        self.energy = energy
        
    def __call__(self, x):
        self.mol.coords = x
        result = numpy.zero(shape=(natoms*3, 1), dtype=REAL_TYPE)
        self.energy.grediant(self.mol, self.ff, result)
        self.mol.forces = -result
        return result

class EngineCallBack:
    def __init__(self, engine, counter = 0):
        self.engine = engine
        self.counter = counter
        
    def __call__(self, x):
        self.counter += 1

class Optimizer(SimulationEngine):
    def __init__(self):
        SimulationEngine.__init__(self)

    @staticmethod
    def add_noise(mol, a = 1.0):
        noise = numpy.random.rand(len(mol.atoms) * 3)
        mol.coords += (noise - 0.5) * a
        
    def optimze(self, mol, ff, energy):
        mol.coords = collect_coords(mol)
        ff.k = collect_k(ff)
        return self.run_optimze(mol, ff, energy)
    
    def run_optimze(self, mol, ff, energy):
        raise NotImplementedError()
         
class DownhillSimplexOptimizer(Optimizer):
    def run_optimze(self, mol, ff, energy, *args, **kwargs):
        efunc = EnergyFunction(mol, ff, energy)
        caller = None # EngineCallBack(self)
        return scipy.optimize.fmin(efunc, mol.coords, callback = caller, *args, **kwargs)
        
                
class PowellOptimizer(Optimizer):
    def run_optimze(self, mol, ff, energy, *args, **kwargs):
        efunc = EnergyFunction(mol, ff, energy)
        caller = None # EngineCallBack(self)
        return scipy.optimize.fmin_powell(efunc, mol.coords, callback = caller, *args, **kwargs)
        
                
